{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# type: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# torchtune Supervised Fine-Tuning\n",
    "\n",
    "This recipe allows TensorZero users to fine-tune models using [torchtune](https://docs.pytorch.org/torchtune/main/) and their own data.\n",
    "Since TensorZero automatically logs all inferences and feedback, it is straightforward to fine-tune a model using your own data and any prompt you want.\n",
    "\n",
    "We demonstrate how to deploy a LoRA fine-tuned model for serverless inference using [Fireworks](https://fireworks.ai). Full instructions to deploy LoRA or full fine-tuned models are provided by [Fireworks](https://docs.fireworks.ai/fine-tuning/fine-tuning-models), [Together](https://docs.together.ai/docs/deploying-a-fine-tuned-model), and other inference providers. You can also use [vLLM](https://docs.vllm.ai/en/latest/examples/online_serving/api_client.html) to serve your fine-tuned model locally. The TensorZero client seemlessly integrates inference using your fine-tuned model for any of these approaches.\n",
    "\n",
    "To get started:\n",
    "\n",
    "- Set your `TENSORZERO_CLICKHOUSE_URL` enironment variable to point to the database containing the historical inferences you'd like to train on.\n",
    "- Set your `HF_TOKEN` to use Llama or Gemma models downloaded through huggingface.\n",
    "- Set the environment variable `CHECKPOINT_HOME` to a path with sufficient storage to save the base LLM checkpoints.\n",
    "- You'll also need to [install](https://docs.fireworks.ai/tools-sdks/firectl/firectl) the CLI tool `firectl` on your machine and sign in with `firectl signin`. You can test that this all worked with `firectl whoami`. We use `firectl` for deployment to Fireworks in this example but you can serve the model however you prefer.\n",
    "- Update the following parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG_PATH = \"../../../examples/data-extraction-ner/config/tensorzero.toml\"\n",
    "\n",
    "FUNCTION_NAME = \"extract_entities\"\n",
    "\n",
    "METRIC_NAME = \"jaccard_similarity\"\n",
    "\n",
    "# The name of the variant to use to grab the templates used for fine-tuning\n",
    "TEMPLATE_VARIANT_NAME = \"gpt_4o_mini\"  # It's OK that this variant uses a different model than the one we're fine-tuning\n",
    "\n",
    "# If the metric is a float metric, you can set the threshold to filter the data\n",
    "FLOAT_METRIC_THRESHOLD = 0.5\n",
    "\n",
    "# Fraction of the data to use for validation\n",
    "VAL_FRACTION = 0.2\n",
    "\n",
    "# Maximum number of samples to use for fine-tuning\n",
    "MAX_SAMPLES = 100_000\n",
    "\n",
    "# Random seed\n",
    "SEED = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select a [supported model](https://docs.pytorch.org/torchtune/main/api_ref_models.html) to fine tune.\n",
    "\n",
    "\n",
    "If you want to use a supported model that we have not yet included in the MODELS dictionary in utils.py, you will need to update the dictionary.\n",
    "\n",
    "Alternatively, you can write the `TUNING_CONFIG` defined below from scratch following the [example configs](https://github.com/pytorch/torchtune/tree/main/recipes/configs) provided by torchtune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The name of the model to fine-tune.\n",
    "MODEL_NAME = \"meta-llama/Meta-Llama-3.1-8B-Instruct\"\n",
    "\n",
    "# Whether to use LoRA or not. Set to False for full model fine-tuning\n",
    "# If set to False, SEVERLESS must also be False as you will need to create your own deployment\n",
    "USE_LORA = True\n",
    "\n",
    "# Whether to use a serverless deployment.\n",
    "# Set to False is full model fine tuning or using LoRA for a model without serverless support\n",
    "SERVERLESS = True\n",
    "\n",
    "# Set to true if you want to include system and user messages in loss calculation\n",
    "TRAIN_ON_INPUT = False\n",
    "\n",
    "# Number of server nodes to use\n",
    "NNODES = 1\n",
    "\n",
    "# Number of devices (e.g., GPUs) to use per node\n",
    "NPROC_PER_NODE = 1\n",
    "\n",
    "# Set the directory where you would like to save the fine-tuned model\n",
    "OUTPUT_DIR = \"fine-tuned\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "tensorzero_path = os.path.abspath(os.path.join(os.getcwd(), \"../../\"))\n",
    "if tensorzero_path not in sys.path:\n",
    "    sys.path.append(tensorzero_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "from utils import MODELS, list_checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert \"CHECKPOINT_HOME\" in os.environ, \"CHECKPOINT_HOME environment variable not set\"\n",
    "assert \"HF_TOKEN\" in os.environ, \"HF_TOKEN environment variable not set\"\n",
    "\n",
    "checkpoint_home = Path(os.environ[\"CHECKPOINT_HOME\"])\n",
    "checkpoint_dir = checkpoint_home / MODEL_NAME\n",
    "\n",
    "command = [\n",
    "    \"tune\",\n",
    "    \"download\",\n",
    "    MODEL_NAME,\n",
    "    \"--output-dir\",\n",
    "    f\"{checkpoint_dir}\",\n",
    "]\n",
    "print(\" \".join(command))\n",
    "try:\n",
    "    subprocess.run(command, check=True)\n",
    "except subprocess.CalledProcessError as e:\n",
    "    print(\"Error occurred:\", e.stderr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally, use Low Rank Adaptation.\n",
    "\n",
    "Some [Fireworks Models]() support [serverless LoRA deployment](https://docs.fireworks.ai/fine-tuning/fine-tuning-models), but full fine-tuning usually needs some form of reserved capacity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if USE_LORA:\n",
    "    MODEL_CONFIG = {\n",
    "        \"_component_\": MODELS[MODEL_NAME][\"modules\"][\"lora\"],\n",
    "        \"lora_attn_modules\": [\n",
    "            \"q_proj\",\n",
    "            \"k_proj\",\n",
    "            \"v_proj\",\n",
    "            \"output_proj\",\n",
    "        ],\n",
    "        \"apply_lora_to_mlp\": True,\n",
    "        \"apply_lora_to_output\": False,\n",
    "        \"lora_rank\": 8,  # higher increases accuracy and memory\n",
    "        \"lora_alpha\": 16,  # usually alpha=2*rank\n",
    "        \"lora_dropout\": 0.0,\n",
    "    }\n",
    "else:\n",
    "    MODEL_CONFIG = {\n",
    "        \"_component_\": MODELS[MODEL_NAME][\"modules\"][\"full\"],\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the training parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOKENIZER_CONFIG = MODELS[MODEL_NAME][\"tokenizer\"]\n",
    "if TOKENIZER_CONFIG.get(\"path\"):\n",
    "    TOKENIZER_CONFIG[\"path\"] = str(checkpoint_dir / TOKENIZER_CONFIG[\"path\"])\n",
    "if TOKENIZER_CONFIG.get(\"merges_file\"):\n",
    "    TOKENIZER_CONFIG[\"merges_file\"] = str(checkpoint_dir / TOKENIZER_CONFIG[\"merges_file\"])\n",
    "TOKENIZER_CONFIG[\"max_seq_len\"] = None  # Can set to an integer value to reduce your memory footprint\n",
    "\n",
    "TUNING_CONFIG = {\n",
    "    \"output_dir\": OUTPUT_DIR,\n",
    "    # Tokenizer\n",
    "    \"tokenizer\": TOKENIZER_CONFIG,\n",
    "    # Model Arguments\n",
    "    \"model\": MODEL_CONFIG,\n",
    "    \"checkpointer\": {\n",
    "        \"_component_\": \"torchtune.training.FullModelHFCheckpointer\",\n",
    "        \"checkpoint_dir\": str(checkpoint_dir),\n",
    "        \"checkpoint_files\": list_checkpoints(checkpoint_dir),\n",
    "        \"recipe_checkpoint\": None,\n",
    "        \"output_dir\": OUTPUT_DIR,\n",
    "        \"model_type\": MODELS[MODEL_NAME][\"checkpointer\"][\"model_type\"],\n",
    "    },\n",
    "    \"resume_from_checkpoint\": False,\n",
    "    \"save_adapter_weights_only\": USE_LORA,\n",
    "    # Optimizer and Scheduler\n",
    "    \"optimizer\": {\n",
    "        \"_component_\": \"torch.optim.AdamW\",\n",
    "        \"fused\": True,\n",
    "        \"weight_decay\": 0.01,\n",
    "        \"lr\": 1e-4,\n",
    "    },\n",
    "    \"lr_scheduler\": {\n",
    "        \"_component_\": \"torchtune.training.lr_schedulers.get_cosine_schedule_with_warmup\",\n",
    "        \"num_warmup_steps\": 100,\n",
    "    },\n",
    "    \"loss\": {\n",
    "        \"_component_\": \"torchtune.modules.loss.LinearCrossEntropyLoss\",\n",
    "    },\n",
    "    # Training\n",
    "    \"epochs\": 1,\n",
    "    \"batch_size\": 2,\n",
    "    \"batch_size_val\": 2,\n",
    "    \"max_steps_per_epoch\": None,\n",
    "    \"gradient_accumulation_steps\": 8,  # Use to increase effective batch size\n",
    "    \"clip_grad_norm\": None,\n",
    "    \"compile\": False,  # torch.compile the model + loss: True increases speed + decreases memory\n",
    "    \"run_val_every_n_steps\": 10,\n",
    "    \"seed\": SEED,\n",
    "    \"shuffle\": True,\n",
    "    # Logging\n",
    "    \"log_every_n_steps\": 1,\n",
    "    \"log_peak_memory_stats\": True,\n",
    "    \"log_level\": \"INFO\",  # DEBUG, WARN, etc.\n",
    "    # Environment\n",
    "    \"device\": \"cuda\",\n",
    "    \"dtype\": \"bf16\",\n",
    "    \"enable_activation_checkpointing\": True,  # True reduces memory\n",
    "    \"enable_activation_offloading\": True,  # True reduces memory\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import tempfile\n",
    "\n",
    "import toml\n",
    "import yaml\n",
    "from tensorzero import (\n",
    "    FloatMetricFilter,\n",
    "    TensorZeroGateway,\n",
    ")\n",
    "from tensorzero.util import uuid7\n",
    "from util import tensorzero_rendered_samples_to_conversations, train_val_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load and render the stored inferences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorzero_client = TensorZeroGateway.build_embedded(\n",
    "    config_file=CONFIG_PATH,\n",
    "    clickhouse_url=os.environ[\"TENSORZERO_CLICKHOUSE_URL\"],\n",
    "    timeout=15,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the metric filter as needed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comparison_operator = \">=\"\n",
    "metric_node = FloatMetricFilter(\n",
    "    metric_name=METRIC_NAME,\n",
    "    value=FLOAT_METRIC_THRESHOLD,\n",
    "    comparison_operator=comparison_operator,\n",
    ")\n",
    "# from tensorzero import BooleanMetricFilter\n",
    "# metric_node = BooleanMetricFilter(\n",
    "#     metric_name=METRIC_NAME,\n",
    "#     value=True  # or False\n",
    "# )\n",
    "\n",
    "metric_node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Query the inferences and feedback from ClickHouse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stored_inferences = tensorzero_client.experimental_list_inferences(\n",
    "    function_name=FUNCTION_NAME,\n",
    "    variant_name=None,\n",
    "    output_source=\"inference\",  # could also be \"demonstration\"\n",
    "    filters=metric_node,\n",
    "    limit=MAX_SAMPLES,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Render the stored inferences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rendered_samples = tensorzero_client.experimental_render_samples(\n",
    "    stored_samples=stored_inferences,\n",
    "    variants={FUNCTION_NAME: TEMPLATE_VARIANT_NAME},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split the data into training and validation sets for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_samples, eval_samples = train_val_split(\n",
    "    rendered_samples,\n",
    "    val_size=VAL_FRACTION,\n",
    "    last_inference_only=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert the rendered samples to conversations for tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_conversations = tensorzero_rendered_samples_to_conversations(train_samples, conversation_key=\"messages\")\n",
    "eval_conversations = tensorzero_rendered_samples_to_conversations(eval_samples, conversation_key=\"messages\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as temp_dir:\n",
    "    temp_dir = Path(temp_dir)\n",
    "\n",
    "    # Write training JSONL\n",
    "    train_json_path = temp_dir / \"train.json\"\n",
    "    with train_json_path.open(\"w\") as f:\n",
    "        for item in train_conversations:\n",
    "            json.dump(item, f)\n",
    "            f.write(\"\\n\")\n",
    "\n",
    "    # Write evaluation JSONL\n",
    "    val_json_path = temp_dir / \"eval.json\"\n",
    "    with val_json_path.open(\"w\") as f:\n",
    "        for item in eval_conversations:\n",
    "            json.dump(item, f)\n",
    "            f.write(\"\\n\")\n",
    "\n",
    "    # Write YAML config\n",
    "    config_path = temp_dir / \"custom_8B_lora_single_device.yaml\"\n",
    "    TUNING_CONFIG[\"dataset\"] = {\n",
    "        \"_component_\": \"torchtune.datasets.chat_dataset\",\n",
    "        \"source\": \"json\",\n",
    "        \"packed\": False,  # True increases speed\n",
    "        \"data_files\": str(train_json_path),\n",
    "        \"conversation_column\": \"messages\",\n",
    "        \"conversation_style\": \"openai\",\n",
    "        \"train_on_input\": TRAIN_ON_INPUT,\n",
    "    }\n",
    "    TUNING_CONFIG[\"dataset_val\"] = {\n",
    "        \"_component_\": \"torchtune.datasets.chat_dataset\",\n",
    "        \"source\": \"json\",\n",
    "        \"packed\": False,  # True increases speed\n",
    "        \"data_files\": str(val_json_path),\n",
    "        \"conversation_column\": \"messages\",\n",
    "        \"conversation_style\": \"openai\",\n",
    "        \"train_on_input\": TRAIN_ON_INPUT,\n",
    "    }\n",
    "    TUNING_CONFIG[\"metric_logger\"] = {\n",
    "        \"_component_\": \"torchtune.training.metric_logging.DiskLogger\",\n",
    "        \"log_dir\": str(temp_dir / \"logs\"),\n",
    "    }\n",
    "    TUNING_CONFIG[\"profiler\"] = {  # Disabled\n",
    "        \"_component_\": \"torchtune.training.setup_torch_profiler\",\n",
    "        \"enabled\": False,\n",
    "        \"output_dir\": str(temp_dir / \"profiling_outputs\"),\n",
    "        \"cpu\": True,\n",
    "        \"cuda\": True,\n",
    "        \"profile_memory\": False,\n",
    "        \"with_stack\": False,\n",
    "        \"record_shapes\": True,\n",
    "        \"with_flops\": False,\n",
    "        \"wait_steps\": 5,\n",
    "        \"warmup_steps\": 3,\n",
    "        \"active_steps\": 2,\n",
    "        \"num_cycles\": 1,\n",
    "    }\n",
    "    with open(config_path, \"w\") as fp:\n",
    "        yaml.safe_dump(\n",
    "            TUNING_CONFIG,\n",
    "            fp,\n",
    "            sort_keys=False,\n",
    "            default_flow_style=False,  # expand lists/dicts in block style\n",
    "        )\n",
    "    print(f\"Config written to {config_path}\")\n",
    "    command = [\n",
    "        \"tune\",\n",
    "        \"run\",\n",
    "        \"--nnodes\",\n",
    "        str(NNODES),\n",
    "        \"--nproc_per_node\",\n",
    "        str(NPROC_PER_NODE),\n",
    "        \"lora_finetune_distributed\" if USE_LORA else \"full_finetune_distributed\",\n",
    "        \"--config\",\n",
    "        str(config_path),\n",
    "    ]\n",
    "    try:\n",
    "        subprocess.run(command, check=True)\n",
    "    except subprocess.CalledProcessError as e:\n",
    "        print(\"Error occurred:\", e.stderr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the model is done training, we need to [deploy](https://docs.fireworks.ai/fine-tuning/fine-tuning-models#deploying-and-using-a-model) it to Fireworks serverless inference. If you need high or guaranteed throughput you can also deploy the model to [reserved capacity](https://docs.fireworks.ai/deployments/reservations) or an on-demand [deployment](https://docs.fireworks.ai/guides/ondemand-deployments)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model_id = \"llama-v3p1-8b-instruct\"\n",
    "base_model_path = f\"accounts/fireworks/models/{base_model_id}\"\n",
    "\n",
    "fine_tuned_model_id = f\"{MODEL_NAME.lower().replace('/', '-').replace('.', 'p')}-{str(uuid7()).split('-')[-1]}\"\n",
    "\n",
    "checkpoint_dir = Path(OUTPUT_DIR) / f\"epoch_{TUNING_CONFIG['epochs'] - 1}\"\n",
    "\n",
    "command = [\n",
    "    \"firectl\",\n",
    "    \"create\",\n",
    "    \"model\",\n",
    "    fine_tuned_model_id,\n",
    "    str(checkpoint_dir),\n",
    "    \"--base-model\",\n",
    "    base_model_path,\n",
    "]\n",
    "try:\n",
    "    result = subprocess.run(command, capture_output=True)\n",
    "    stdout = result.stdout.decode(\"utf-8\")\n",
    "    print(\"Command output:\", stdout)\n",
    "except subprocess.CalledProcessError as e:\n",
    "    print(\"Error occurred:\", e.stderr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model_id(stdout: str) -> str:\n",
    "    for line in stdout.splitlines():\n",
    "        if line.strip().startswith(\"Name:\"):\n",
    "            return line.split(\":\")[1].strip()\n",
    "    raise ValueError(\"Model ID not found in output\")\n",
    "\n",
    "\n",
    "model_identifier = get_model_id(stdout)\n",
    "\n",
    "model_identifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a deployment if not using a model with serverless support, if it does not support serveless addons, or if you are doing full fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not SERVERLESS:\n",
    "    command = [\"firectl\", \"create\", \"deployment\", model_identifier]\n",
    "    print(\" \".join(command))\n",
    "    result = subprocess.run(command, capture_output=True)\n",
    "    if result.returncode != 0:\n",
    "        print(result.stderr.decode(\"utf-8\"))\n",
    "    else:\n",
    "        stdout = result.stdout.decode(\"utf-8\")\n",
    "        print(stdout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the LoRA addon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if USE_LORA:\n",
    "    command = [\"firectl\", \"load-lora\", model_identifier]\n",
    "    print(\" \".join(command))\n",
    "    result = subprocess.run(command, capture_output=True)\n",
    "    if result.returncode != 0:\n",
    "        print(result.stderr.decode(\"utf-8\"))\n",
    "    else:\n",
    "        stdout = result.stdout.decode(\"utf-8\")\n",
    "        print(stdout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the model is deployed, you can add the fine-tuned model to your config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_config = {\n",
    "    \"models\": {\n",
    "        model_identifier: {\n",
    "            \"routing\": [\"fireworks\"],\n",
    "            \"providers\": {\"fireworks\": {\"type\": \"fireworks\", \"model_name\": model_identifier}},\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "print(toml.dumps(model_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You're all set!\n",
    "\n",
    "You can change the weight to enable a gradual rollout of the new model."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "ipynb,py:percent",
   "main_language": "python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
